/*
 * Project CelestialBeing: allocators.hpp
 * Created: 2023/12/08
 * This program uses MulanPSL2. See LICENSE for more.
 */

#ifndef CELESTIALBEING_ALLOCATORS_HPP
#define CELESTIALBEING_ALLOCATORS_HPP
#include <vector>
#include <list>
#include "cb/common/celestial_being_base.hpp"
#include "base_math.hpp"

CB_NAMESPACE_BEGIN

// The simple linear allocator for fixed size blocks on the heap.
// NO THREAD SAFETY GUARANTEES.
template<size_t ElementSize, size_t NumBlockElements, size_t Alignment>
class TFixedSizeBlockHeapAllocator {
public:
    static_assert(ElementSize > 0, "ElementSize should be greater than 0.");
    static_assert(NumBlockElements > 0, "NumBlockElements should be greater than 0.");
    static_assert(Alignment > 0, "Alignment should be greater than 0.");
    static_assert(IsPowerOf2(Alignment), "Alignment should be a power of 2.");

    TFixedSizeBlockHeapAllocator() = default;

    // Allocate an element from one of the blocks with free slots.
    // If there is no block with free slots, allocate a new block.
    inline void * Allocate () {
        Block * block;
        if (blocks_with_free_slots_.empty())
            block = AllocateBlock();
        else block = blocks_with_free_slots_.back();
        void * element = block->data + ElementSize * block->free_slots[--block->num_free_slots];
        if(!block->num_free_slots) blocks_with_free_slots_.pop_back();
        return element;
    }

    // Free an element.
    // The element must be allocated by this allocator.
    // The element must not be freed before.
    inline void Free (void * element) {
        // Just iterate over all the blocks.
        // TODO speed up?
        for (auto block : blocks_with_free_slots_) {
            if (element >= block->data && element < block->data + ElementSize * NumBlockElements) {
                assert(block->num_free_slots < NumBlockElements);
                block->free_slots[block->num_free_slots++] = (std::byte *)element - block->data;
                if(block->num_free_slots == 1) blocks_with_free_slots_.emplace_back(block);
                return;
            }
        }
        assert(false && "The element is not allocated by this allocator or has been freed before.");
    }

    inline ~TFixedSizeBlockHeapAllocator () {
        assert(blocks_with_free_slots_.size() == num_blocks_ && "Some blocks are not freed.");
        Block * block = block_head_;
        while (block) {
            assert(block->num_free_slots == NumBlockElements && "Some elements are not freed."");
            Block * next_block = block->next_block;
            ::operator delete(block, std::align_val_t(Alignment));
            block = next_block;
        }
    }

private:
    struct Block {
        std::byte data[RoundUpToMultiplyOf(ElementSize * NumBlockElements, Alignment)];
        uint32_t free_slots[NumBlockElements];
        std::byte num_free_slots{};
        Block * prev_block{};
        Block * next_block{};
        Block (Block* next): next_block(next) {}
    };

    // Allocate a block of memory.
    // Return nullptr if the allocation fails.
    inline void *AllocateBlock () {
        void * block_addr = ::operator new(sizeof(Block), std::align_val_t(Alignment));
        block_head_->prev_block = new(block_addr) Block(block_head_);
        block_head_ = block_head_->prev_block;
        blocks_with_free_slots_.emplace_back(block_head_);
        num_blocks_ ++;
        return block_head_;
    }

    Block * block_head_;
    std::list<Block *> blocks_with_free_slots_;
    size_t num_blocks_{};

};

CB_NAMESPACE_END

#endif //CELESTIALBEING_ALLOCATORS_HPP
